# ------------------------------------------------------------------------------
# YOLO Object Detection Detector
# ------------------------------------------------------------------------------
# 使用 YOLO 进行物体检测
# ------------------------------------------------------------------------------

import cv2
from .base_detector import BaseDetector

try:
    from ultralytics import YOLO
    YOLO_AVAILABLE = True
except ImportError:
    YOLO_AVAILABLE = False


class YOLODetector(BaseDetector):
    """YOLO 物体检测器"""
    
    def __init__(self, model_name="yolo11n.pt"):
        super().__init__()
        
        if not YOLO_AVAILABLE:
            raise ImportError("YOLO 不可用。请安装: pip install ultralytics")
        
        try:
            # 加载 YOLO 模型
            self.model = YOLO(model_name)
            self.model_loaded = True
        except Exception as e:
            print(f"加载 YOLO 模型失败: {e}")
            self.model_loaded = False
            
    def process_frame(self, frame):
        """
        使用 YOLO 检测物体
        
        Args:
            frame: 输入图像帧
            
        Returns:
            annotated_frame: 标注了检测结果的图像帧
        """
        self.update_fps()
        
        if not self.model_loaded:
            cv2.putText(frame, 'YOLO 模型未加载', (50, 50), 
                       cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)
            frame = self.visualize_fps(frame)
            return frame
        
        try:
            # 运行 YOLO 推理
            results = self.model(frame, verbose=False)
            
            # 绘制检测结果
            annotated_frame = results[0].plot()
            
            # 显示检测到的物体数量
            num_detections = len(results[0].boxes)
            detection_text = f'检测到 {num_detections} 个物体'
            cv2.putText(annotated_frame, detection_text, (24, 50), 
                       cv2.FONT_HERSHEY_PLAIN, 1, (0, 255, 0), 1)
            
            annotated_frame = self.visualize_fps(annotated_frame)
            return annotated_frame
            
        except Exception as e:
            print(f"YOLO 推理错误: {e}")
            cv2.putText(frame, f'错误: {str(e)}', (50, 50), 
                       cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2)
            frame = self.visualize_fps(frame)
            return frame
    
    def cleanup(self):
        """清理资源"""
        if hasattr(self, 'model'):
            del self.model

